﻿module TypeRules

open ExceptionMonad
open TinyJava.Ast
open Environment

type State<'a> = Environment -> ('a * Environment) StatementOrException
type TypeCheckState = State<TinyJava.Ast.SynType>

let state = StateBuilder()

let return_error = "Not all control paths return a value"

//check if two types are compatible
let check_type (t1 : SynType) (t2 : SynType) (is_assignment : bool) : State<Unit> =
  state
    {
      match (t1,t2) with
      | SynType.Int,SynType.Int
      | SynType.Boolean,SynType.Boolean
      | SynType.String,SynType.String -> return ()
      | SynType.Array(t1),SynType.Array(t2) when t1 = t2 -> return ()
      | SynType.Identifier(_),SynType.Identifier("**TjNullVal") -> return()
      | SynType.Array(t1),t2 when t1 = t2 && is_assignment -> return ()
      | SynType.Identifier(c1_name),SynType.Identifier(c2_name) when c1_name = c2_name -> return ()
      | SynType.Identifier(c1_name),SynType.Identifier(c2_name) ->
          let! env = get_state
          if is_super_class env.SymbolTable c1_name c2_name |> not then
            return! fail "TYPE MISMATCH: The assigned class variable is not super class of the right value"
          else
            return ()
      | SynType.Int,_ -> return! fail "TYPE MISMATCH: Integer expected"
      | SynType.Boolean,_ -> return! fail "TYPE MISMATCH: Boolean expected"
      | SynType.String,_ -> return! fail "TYPE MISMATCH: String expected"
      | SynType.Array(t1),t2 when t1 <> t2 && is_assignment -> return! fail "TYPE MISMATCH: Array type does not match assignment"
      | SynType.Array(_),_ -> return! fail "TYPE MISMATCH: Array expected"
      | SynType.Identifier(_),_ -> return! fail "TYPE MISMATCH: Class expected"
    }

let rec eval_program (p : Program) : State<Unit> =
  state
    {
      for class_decl in p do
        do! eval_class_decl class_decl
    }

and eval_class_decl (c : ClassDecl) : State<Unit> =
  state
    {
      let eval_class_body (fm_list : FieldOrMethodDecl list) =
        state
          {
            for fm in fm_list do
              do! eval_field_or_method fm
          }
      let! env = get_state   
      let (current_id,parent_class_id,fm_list) = c
      match parent_class_id with
      |Some parent_class ->
        let pc = try_find_class env.SymbolTable parent_class
        match pc with
        | Some (Identifier(parent_id)) ->
            let new_env = {env with CurrentClass = current_id} 
            do! set_state new_env
            do! eval_class_body fm_list
        | None -> return! fail "Parent class definition missing"
      | None ->
          let new_env = {env with CurrentClass = current_id} 
          do! set_state new_env
          do! eval_class_body fm_list
    }

and eval_field_or_method (fm : FieldOrMethodDecl) : State<Unit> =
  state
    {
      let! env = get_state
      match fm with
      | FieldDecl(t,id_list) -> return ()
      | MethodDecl(ret_type,id,formals,block) ->
          let block_table =
            {
              Vars = []
              Parent = None
              Owner = Some(get_method_info env.SymbolTable (env.CurrentClass) id)
              Child = None
            }
          env.SymbolTable.[env.CurrentClass].Methods.[id].Body := Some (block_table)
          do! set_state {env with BlockTable = ref(block_table)}
          let! env = get_state
          match ret_type with
          | RetType(_) -> do! set_state {env with CheckReturn = true;ReturnFound = false}
          | Void -> do! set_state {env with CheckReturn = false;ReturnFound = false}
          let! env = get_state
          do! eval_block block
          let! env = get_state
          if env.ReturnFound |> not && env.CheckReturn then
            return! fail return_error
          else
            return ()
    }

and eval_block (block : Block) : State<Unit> =
  state
    {
      let (var_decls,statements) = block
      let! env = get_state
      for (t,defs) in var_decls do
        for (name,init) in defs do
          let opt_var_name = try_find_id (BlockTable(!env.BlockTable)) name true env
          match opt_var_name with
          | Some _ -> return! fail "Identifier already defined"
          | None -> 
              match init with
              | Some init_expr ->
                  let! e = eval_expression init_expr
                  do! check_type t e false
                  env.BlockTable.Value.addVar (name,t)
                  return ()
              | None ->
                  env.BlockTable.Value.addVar (name,t)
                  return ()
      for s in statements do
        do! eval_statement s
    }

and eval_literal (l : Literal) : State<SynType> =
  state
    {
      match l with
      | IntLit _ -> return SynType.Int
      | StringLit _ -> return SynType.String
      | Boolean _ -> return SynType.Boolean
      | Null -> return (SynType.Identifier "**TjNullVal")
    }

and eval_location (l : Location) : State<SynType> =
  state
    {
      match l with
      | Id(name) ->
          let! env = get_state
          let found_id = try_find_id (BlockTable(!env.BlockTable)) name false env
          match found_id with
          | Some id_type -> return id_type
          | None -> return! fail ("IDENTIFIER " + name + " NOT FOUND")
      | ClassLoc(e,field_name) ->
          let! env = get_state
          let! e = eval_expression e
          match e with
          | Identifier(c_name) ->
              let field_type = try_find_field env.SymbolTable c_name field_name
              match field_type with
              | Some t -> return t
              | None -> return! fail ("CLASS FIELD " + field_name + " NOT FOUND")
          | _ -> return! fail ("TYPE MISMATCH: Class definition expected")
      | ArrayLoc(e1,e2) ->
          let! e1 = eval_expression e1
          let! e2 = eval_expression e2
          match e1,e2 with
          | SynType.Array(t),SynType.Int -> return t
          | SynType.Array(_),e2 -> return! fail ("TYPE MISMATCH: Array index must be an integer value")
          | e1,SynType.Int -> return! fail ("TYPE MISMATCH: Expected array type")
          | _ -> return! fail ("TYPE MISMATCH: Expected array type or Array index must be an integer value")
    }

and eval_statement (s : Statement) : State<Unit> =
  state
    {
      let! env = get_state
      if env.ReturnFound then
        return! fail "Unreachable statement"
      else
        //if we get to this point ReturnFound is surely false
        match s with
        | Assign(l,e) ->
            let! l = eval_location l
            let! e = eval_expression e
            do! check_type l e true

        | Statement.Call(m,actuals) ->
            let! m,formals = eval_method m
            match actuals with
            | Some actuals ->
                do! eval_actuals actuals formals
                match m with
                | RetType(_) -> return! fail "A statement cannot return a value"
                | Void -> return ()
            | None ->
                match m with
                | RetType(_) -> return! fail "A statement cannot return a value"
                | Void -> return () 

        | Return(e_opt) ->
            let! env = get_state

            //a RETURN statement always grants that this path will return. Update the flag
            do! set_state {env with ReturnFound = true}
            let! env = get_state
            let m_ret = get_method_ret_type !env.BlockTable
            match e_opt with
            | Some e ->
                let! e = eval_expression e
                match m_ret with
                | RetType(t) ->
                    if t = e then
                      return ()
                    else
                      return! fail "TYPE MISMATCH: return expression does not match expected return type"
                | Void ->
                    return! fail "TYPE MISMATCH: VOID does not match return type"
            | None ->
              match m_ret with
              | RetType(_) -> return! fail "TYPE MISMATCH: return expression does not match VOID"
              | Void -> return()

        | If(e,s1,s2_opt) ->
            let! e = eval_expression e
            let! env = get_state
            if (e = SynType.Boolean) then
              do! eval_statement s1
              //an IF alone never grants we always reach a return statement, so we reset the flags
              do! set_state {env with ReturnFound = false}
              match s2_opt with
              | Some s2 ->
                  do! eval_statement s2
                  let! env = get_state
                  //if the ELSE path has a return, then this path will always return
                  if env.ReturnFound |> not && env.CheckReturn then
                    return! fail return_error
                  else
                    return ()
              | None -> return ()
            else
              return! fail "TYPE MISMATCH: If condition must be boolean"

        | While(e,s) ->
            let! env = get_state
            let! e = eval_expression e
            if e = SynType.Boolean then
              do! set_state {env with LoopCounter = env.LoopCounter + 1}
              let! env = get_state
              do! eval_statement s
              //a WHILE statement never grants we reach a return statement since its condition might be never met.
              //Reset the flag.
              do! set_state {env with LoopCounter = env.LoopCounter - 1;ReturnFound = false}
            else
              return! fail "TYPE MISMATCH: While condition must be boolean"

        | Block(b) ->
            let! env = get_state
            let new_block_table =
              {
                Vars = []
                Parent = Some(ref !env.BlockTable)
                Owner = env.BlockTable.Value.Owner
                Child = None
              }
            env.BlockTable.Value.Child <- Some(ref new_block_table)
            env.BlockTable := new_block_table
            do! eval_block b

        | Break 
        | Continue ->
            let! env = get_state
            if env.LoopCounter = 0 then
              return! fail "Break or Continue statement not inside a loop"
            else
              return ()
    }

and eval_expression (e : Expression) : State<SynType> =
  state
    {
      match e with
      | Location(loc) -> return! eval_location loc

      | Call(m,actuals) ->
          let! m,formals = eval_method m
          match actuals with
          | Some actuals ->
              do! eval_actuals actuals formals
              match m with
              | RetType(t) -> return t
              | Void -> return! fail "An expression cannot be of type VOID"
          | _ ->
            match m with
            | RetType(t) -> return t
            | Void -> return! fail "An expression cannot be of type VOID"

      | This ->
            let! env = get_state
            return SynType.Identifier(env.CurrentClass)

      | DefaultConstructor(id) ->
          let! env = get_state
          let class_table_opt = try_find_class env.SymbolTable id
          match class_table_opt with
          | Some table -> return Identifier(id)
          | None -> return! fail "Non-existing class definition"

      | ArrayConstructor(t,e) ->
          let! e = eval_expression e
          match e with
          | SynType.Int -> return! eval_array_type t
          | _ -> return! fail "Array index must be an integer value"
      
      | Length(e) ->
          let! e = eval_expression e
          match e with
          | SynType.Array(_) -> return SynType.Int
          | _ -> return! fail "OPERATOR LENGTH can be applied only to an array type"

      | BinaryOp(e1,op,e2) ->
          let! op = eval_binary_op op e1 e2
          return op

      | UnaryOp(op,e) ->
          let! op = eval_unary_op op e
          return op

      | Literal(literal) -> return! eval_literal literal
    }


and eval_array_type (t : SynType) : State<SynType> =
  state
    {
      match t with
      | SynType.Int
      | SynType.Boolean
      | SynType.String -> return SynType.Array(t)
      | SynType.Identifier(class_id) ->
          let! env = get_state
          let class_type_opt = try_find_class env.SymbolTable class_id
          match class_type_opt with
          | Some class_type -> return SynType.Array(class_type)
          | None -> return! fail "Class Definition does not exist"
      | SynType.Array(t1) -> return! eval_array_type t1
    }

and eval_library_function (f_name : Id) : State<MethodRetType * (Formals list)> =
  state
    {
      match f_name with
      | "println" -> return Void,[SynType.String,"s"]
      | "print"   -> return Void,[SynType.String,"s"]
      | "printi"  -> return Void,[SynType.Int,"i"]
      | "printb"  -> return Void,[SynType.Boolean,"b"]
      | "readi"   -> return RetType(SynType.Int),[]
      | "readln"  -> return RetType(SynType.String),[]
      | "eof"     -> return RetType(SynType.Boolean),[]
      | "stoi"    -> return RetType(SynType.Int),[SynType.String,"s";SynType.Int,"n"]
      | "itos"    -> return RetType(SynType.String),[SynType.Int,"i"]
      | "stoa"    -> return RetType(SynType.Array(SynType.Int)),[SynType.String,"s"]
      | "aots"    -> return RetType(SynType.String),[SynType.Array(SynType.Int),"a"]
      | "random"  -> return RetType(SynType.Int),[SynType.Int,"n"]
      | "time"    -> return RetType(SynType.Int),[]
      | "exit"    -> return RetType(SynType.Int),[SynType.Int,"i"]
      | _         -> return! fail "Bug in type checker at eval_library_function"
    }

and eval_method (m : Method) : State<MethodRetType * (Formals list)> =
  state
    {
      let! env = get_state
      match m with
      | ImplicitMethod(m_name) ->
          let m_opt = try_find_method env.SymbolTable env.CurrentClass m_name
          match m_opt with
          | Some (ret_type,formals) -> return (ret_type,formals)
          | None ->
              if is_library_function m_name then
                return! eval_library_function m_name
              else
                return! fail ("UNDEFINED METHOD " + m_name)
      | MethodCall(e,m_name) ->
          let! e = eval_expression e
          match e with
          | Identifier(class_id) ->
              let c = try_find_class env.SymbolTable class_id
              match c with
              | Some (Identifier(c_id)) ->
                  let m_opt = try_find_method env.SymbolTable c_id m_name
                  match m_opt with
                  | Some (ret,formals) -> return (ret,formals)
                  | None -> return! fail "Method definition missing"
              | None -> return! fail "Class definition missing"
              | _ -> return! fail "Bug in type checker at eval_method"
          | _ -> return! fail ("TYPE MISMATCH: Class type expected")
    }

and eval_actuals (actuals : Actuals) (formals : Formals list) : State<Unit> =
  state
    {
      if actuals.Length < formals.Length then
        return! fail "Call values are less than the method parameters"
      elif actuals.Length > formals.Length then
        return! fail "Call values are more than the method parameters"
      else
        let afs = List.zip actuals formals
        match afs with
        | [] -> return ()
        | (a,f)::afs ->
            let! a = eval_expression a
            let (f_type,f_name) = f
            if a = f_type then
              return ()
            else
              return! fail ("Call parameter does not match the type of " + f_name)
    }

and eval_binary_op (op : BinaryOp) (e1: Expression) (e2 : Expression) : State<SynType> =
  state
    {
      match op with
      | Plus(op_t) ->
          let! e1 = eval_expression e1
          let! e2 = eval_expression e2
          match e1,e2 with
          | SynType.Int,SynType.Int ->
              op_t := SynType.Int 
              return SynType.Int
          | SynType.String,SynType.String ->
              op_t := SynType.String
              return SynType.String
          | _ -> return! fail "PLUS OPERATOR ERROR: Expected Integer or string"
      | Minus
      | Mul
      | Div
      | Mod -> 
        let! e1 = eval_expression e1
        let! e2 = eval_expression e2
        match e1,e2 with
        | SynType.Int,SynType.Int -> return SynType.Int
        | _ -> return! fail "ARITHMETIC OPERATOR ERROR: Integer expected"
      | And
      | Or ->
        let! e1 = eval_expression e1
        let! e2 = eval_expression e2
        match e1,e2 with
        | SynType.Boolean,SynType.Boolean -> return SynType.Boolean
        | _ -> return! fail "BOOLEAN OPERATOR ERROR: Boolean expected"
      | Less
      | Leq
      | Greater
      | Geq ->
        let! e1 = eval_expression e1
        let! e2 = eval_expression e2
        match e1,e2 with
        | SynType.Int,SynType.Int -> return SynType.Boolean
        | _ -> return! fail "INTEGER COMPARISON OPERATOR ERROR: Integer expected"
      | Eq
      | Neq ->
        let! e1 = eval_expression e1
        let! e2 = eval_expression e2
        match e1,e2 with
        | SynType.Int,SynType.Int -> return SynType.Boolean
        | SynType.Boolean,SynType.Boolean -> return SynType.Boolean
        | SynType.Identifier(s1),SynType.Identifier(s2) when s1 = s2 -> return SynType.Boolean
        | SynType.String,SynType.String -> return SynType.Boolean
        | SynType.Array(t1),SynType.Array(t2) when t1 = t2-> return SynType.Boolean
        | _ -> return! fail "EQUALITY OPERATOR ERROR: type mismatch"
    }

and eval_unary_op (op : UnaryOp) (e : Expression) : State<SynType> =
  state
    {
      match op with
      | Sign ->
        let! e = eval_expression e
        match e with
        | SynType.Int -> return SynType.Int
        | _ -> return! fail "SIGN OPERATOR ERROR: Integer expected"
      | Not ->
        let! e = eval_expression e
        match e with
        | SynType.Boolean -> return SynType.Boolean
        | _ -> return! fail "NOT OPERATOR ERROR: Boolean expected"
    }
